import os
import threading
import time
from io import BytesIO

import requests
from PIL import Image
import cv2

"""
一些处理数据集的函数；
picknone ：抓取链接图片
pickfromvideo ： 抓取视频中固定帧间隔的图片
process_video ： 平滑抓取视频中N帧图片
"""

# 从链接抓图
def picknone():
    # 设置图片存储路径
    SAVE_PATH = "none/"

    # 设置url链接
    URL = " "

    # 定义函数来抓取并保存图片
    def save_image():
        # 发送请求获取图片
        response = requests.get(URL)

        # 创建图片对象
        img = Image.open(BytesIO(response.content))

        # 检查文件夹是否存在，不存在则创建
        if not os.path.exists(SAVE_PATH):
            os.makedirs(SAVE_PATH)

        # 生成文件名，格式为：年月日-时分秒.jpg
        filename = time.strftime("%Y%m%d-%H%M%S.jpg")

        # 保存图片
        img.save(os.path.join(SAVE_PATH, filename))

    # 每10分钟抓取一次图片
    while True:
        save_image()
        time.sleep(600)  # 10分钟=600秒

# 从视频平均抓取图片,指定帧间隔
def pickfromvideo():

    # 打开视频文件
    cap = cv2.VideoCapture("video/192.168.101.244_01_20230521163601702.mp4")

    # 获取视频信息
    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = cap.get(cv2.CAP_PROP_FPS)

    # 截取10帧，间隔均匀，分布均匀
    interval = int(frame_count / 12)
    frames = []
    for i in range(frame_count):
        ret, frame = cap.read()
        if i % interval == 0:
            frames.append(frame)
            cv2.imwrite(f"test/frame_{i}.jpg", frame)
        if not ret:
            break
    # 释放资源
    cap.release()

# 从视频平均抓取图片,平均取，
# 可将任意视频长度视频平均截取n帧
def process_video(filename,input_folder,output_folder):

    # 生成输出文件夹路径
    output_foldername = os.path.splitext(filename)[0]
    output_folderpath = os.path.join(output_folder, output_foldername)
    os.makedirs(output_folderpath, exist_ok=True)

    # 读取视频文件
    video_path = os.path.join(input_folder, filename)
    video_capture = cv2.VideoCapture(video_path)

    # 计算帧率并获取总帧数
    # fps = video_capture.get(cv2.CAP_PROP_FPS)
    total_frames = int(video_capture.get(cv2.CAP_PROP_FRAME_COUNT))

    """
    Nframe为抓取视频的帧数，12为，平滑抽取视频中12帧图片
    """
    Nframe = 12

    # 生成24帧的等间隔采样列表
    samples = [int(i * total_frames / Nframe) for i in range(Nframe)]

    # 循环读取每个采样帧并保存为JPG图像
    count = 0
    for i, sample in enumerate(samples):
        # 设置视频读取指针位置
        video_capture.set(cv2.CAP_PROP_POS_FRAMES, sample)

        # 读取当前采样帧
        success, image = video_capture.read()

        if success:
            # 生成输出文件路径和文件名
            output_filename = str(count) + ".jpg"
            output_path = os.path.join(output_folderpath, output_filename)

            # 保存当前帧为JPG图像
            cv2.imwrite(output_path, image)

            count += 1

    # 释放视频捕获器
    video_capture.release()


# 使用多线程处理
def thredpick():
    # 遍历输入文件夹中的所有MP4文件
    input_folder = "video/ZF"
    output_folder = "img/ZF"

    for filename in os.listdir(input_folder):
        t = threading.Thread(target=process_video, args=(filename,input_folder,output_folder))
        t.start()

if __name__ == '__main__':
    pickfromvideo()